import { useRef, useState, useCallback, useEffect } from "react";

export type SSEStatus = "idle" | "connecting" | "open" | "closed" | "error";

export interface UseSSEOptions extends EventSourceInit {
  generator?: (
    messageData: string
  ) =>
    | Record<string, string | string[] | number | number[]>
    | null
    | Promise<Record<string, string | string[] | number | number[]> | null>;
}

export function useSSE() {
  const [messages, setMessages] = useState<string[]>([]);
  const [sseData, setSSEData] = useState<Record<string, string | string[] | number | number[]>[]>(
    []
  );
  const [status, setStatus] = useState<SSEStatus>("idle");
  const [error, setError] = useState<Event | null>(null);
  const eventSourceRef = useRef<EventSource | null>(null);

  const connect = useCallback((url: string, options?: UseSSEOptions) => {
    // 关闭已有连接
    if (eventSourceRef.current) {
      eventSourceRef.current.close();
    }
    setMessages([]);
    setSSEData([]);
    setStatus("connecting");
    setError(null);
    const es = new EventSource(url, options);
    eventSourceRef.current = es;

    es.onopen = () => setStatus("open");
    es.onmessage = async (event) => {
      const messageData = event.data as string; // 明确消息数据类型
      setMessages((prev) => [...prev, messageData]);

      if (options?.generator) {
        const data = await options.generator(messageData);
        if (!data) return;
        setSSEData((prev) => [...prev, data]);
      }
    };
    es.addEventListener("end", (e) => {
      setMessages((prev) => [...prev, e.data]);
      es.close();
      setStatus("closed");
    });
    es.onerror = (event) => {
      setStatus("error");
      setError(event);
    };
  }, []);

  const close = useCallback(() => {
    if (eventSourceRef.current) {
      eventSourceRef.current.close();
      eventSourceRef.current = null;
      setStatus("closed");
    }
  }, []);

  // 组件卸载时自动关闭
  useEffect(() => {
    return () => {
      if (eventSourceRef.current) {
        eventSourceRef.current.close();
        eventSourceRef.current = null;
      }
    };
  }, []);

  return { messages, sseData, status, error, connect, close }; // 返回 filteredMessages
}
